#include "DefaultSpellChecker.h"

#include "direct_product.h"
#include "wordDistortionGenerator.h"

#include <algorithm>
#include <string>
#include <vector>
#include <utility>
#include <limits>

#include <cmath>


using std::string;
using std::vector;

namespace NSpellCheck {

static inline double smooth(const int* freq) {
  return freq ? *freq : 1;
}

static inline double smooth(int freq) {
  return freq > 0 ? freq : 1;
}


DefaultSpellChecker::DefaultSpellChecker(Trie<int>* unigrams, BigramCollection* bigrams):
        unigrams(unigrams),
        bigrams(bigrams)
{
}

DefaultSpellChecker::~DefaultSpellChecker() {}

std::pair<size_t, int> DefaultSpellChecker::findSplitIndex(const string& word) const {
    for (size_t i = 3; i + 3 < word.length(); ++i) {
        const int* freq = (*unigrams)[word.substr(0, i)];
        if (freq && *freq >= 1000) {
            return std::make_pair(i, *freq);
        }
    }
    return std::make_pair(string::npos, 0);
}

void DefaultSpellChecker::findSplitVariations(string word, vector<Variation>* variations) const {
/*
    Variation variation;
    variation.weight = 0.5;
    std::pair<size_t, int> splitIndex;
    while ((splitIndex = findSplitIndex(word)).first != string::npos) {
        variation.phrase.push_back(word.substr(0, splitIndex.first));
        word = word.substr(splitIndex.first);
    }
    const int* freq = (*unigrams)[word];
    if (freq && *freq >= 1000 && !variation.phrase.empty()) {
        variation.phrase.push_back(word);
        variations->push_back(variation);
    }
*/
    const int* f = (*unigrams)[word];
    int freq = f ? *f : 0;
    if (freq > 10000) {
      return;
    }
    wordDistortionGenerator distortion;
    std::vector<std::pair<std::string, std::string> > cuts;
    distortion.generateTokenCuts(word, cuts);
    for (size_t i = 0; i < cuts.size(); ++i) {
        const std::string& w1 = cuts[i].first;
        const std::string& w2 = cuts[i].second;
        if (w1.length() <= 2 || w2.length() <= 2) {
            continue;
        }
        const int* freq1 = (*unigrams)[w1];
        const int* freq2 = (*unigrams)[w2];
        if (freq1 == 0 || freq2 == 0 || *freq1 < 1000 || *freq2 < 1000 || *freq1 < freq || *freq2 < freq || bigrams->getFrequency(freq1, freq2) < 100) {
            continue;
        }
        Phrase phrase;
        phrase.push_back(w1);
        phrase.push_back(w2);
        variations->push_back(Variation(exp(-0.3), phrase));
    }
}

bool DefaultSpellChecker::findMergeVariations( const Phrase& phrase, std::vector<Variation>* variations ) const {
    bool foundNew = false;
    for (size_t i = 0; i + 1 < phrase.size(); ++i) {
        const int* freq1 = (*unigrams)[phrase[i]];
        const int* freq2 = (*unigrams)[phrase[i+1]];
        int bigramFreq = bigrams->getFrequency(freq1, freq2);
        if (bigramFreq > 1000)
            continue;
        string w = phrase[i] + phrase[i+1];
        const int* freq = (*unigrams)[w];
        if (freq && *freq > 100) {
            Phrase newPhrase(phrase.begin(), phrase.begin() + i);
            newPhrase.push_back(w);
            newPhrase.insert(newPhrase.end(), phrase.begin() + i + 1, phrase.end());
            variations->push_back(Variation(exp(-0.5), newPhrase));
            foundNew = true;
        }
    }
    return foundNew;
}

void DefaultSpellChecker::process(const Phrase& phrase, std::vector<Variation>* suggestions) const {
    vector<vector<Variation> > wordVariations(phrase.size());
    size_t maxWordVariations = (size_t)pow(5e6, 1.0 / phrase.size());
    for (size_t i = 0; i < phrase.size(); ++i) {
        findWordVariations(phrase[i], &wordVariations[i]);
        if (wordVariations[i].size() > maxWordVariations) {
            wordVariations[i].resize(maxWordVariations);
        }
    }
    directProduct(wordVariations, suggestions);
    //findMergeVariations(phrase, suggestions);

    for (vector<Variation>::iterator i = suggestions->begin(); i != suggestions->end(); ++i) {
        i->weight *= languageModel(i->phrase);
    }

    size_t maxCount = 6;
    if (suggestions->size() > maxCount) {
        std::nth_element(suggestions->begin(), suggestions->begin() + maxCount, suggestions->end(),
                ReverseVariationComparator());
        suggestions->resize(maxCount);
    }

    std::sort(suggestions->begin(), suggestions->end(), ReverseVariationComparator());

    size_t variationsToKeep = 1;
    // cut off when there is a huge drop of probability
    while (variationsToKeep < suggestions->size() &&
           (*suggestions)[variationsToKeep].weight * 3 >= (*suggestions)[variationsToKeep-1].weight)
    {
        ++variationsToKeep;
    }
    if (suggestions->size() > variationsToKeep) {
        suggestions->resize(variationsToKeep);
    }
}

void DefaultSpellChecker::findWordVariations(const string& word, vector<Variation>* variations) const {
    variations->clear();
    vector<string> edits;
    vector<double> weights;
    wordDistortionGenerator distortion;
    distortion.generateWeightedTwoEditOperatonLocalityDistortion(word, edits, weights, true);
    for (size_t i = 0; i < edits.size(); ++i) {
        const int* freq = (*unigrams)[edits[i]];
        if (freq != 0 && *freq >= 100) {
            Phrase phrase;
            phrase.push_back(edits[i]);
            double correctionWeight = exp(-20.0 * weights[i] / word.length());
            variations->push_back(Variation(correctionWeight, phrase));
        }
    }
    //findSplitVariations(word, variations);
    if (variations->empty()) {
        // fall back to baseline solution
        Phrase phrase;
        phrase.push_back(word);
        variations->push_back(Variation(1.0, phrase));
    }

    std::sort(variations->begin(), variations->end(), ReverseVariationComparator());
    size_t variationsToKeep = 1;
    // cut off when there is a huge drop of probability
    while (variationsToKeep < variations->size() &&
           (*variations)[variationsToKeep].weight * 10 >= (*variations)[variationsToKeep-1].weight)
    {
        ++variationsToKeep;
    }
    if (variations->size() > variationsToKeep) {
        variations->resize(variationsToKeep);
    }
}

double DefaultSpellChecker::languageModel(const Phrase& phrase) const {
    double res = 0.0;
    vector<const int*> iters;
    for (Phrase::const_iterator i = phrase.begin(); i != phrase.end(); ++i) {
        iters.push_back((*unigrams)[*i]);
    }
    for (size_t i = 0; i < iters.size(); ++i) {
        res += log(smooth(iters[i]));
    }
    res = exp(res / phrase.size());

    int numBigrams = 0;
    for (size_t i = 0; i + 1 < iters.size(); ++i) {
        int freqBigram = bigrams->getFrequency(iters[i], iters[i+1]);
        if (freqBigram > 0) {
            ++numBigrams;
        }
    }

    res *= exp((double)numBigrams);
    return res;
}

} //ns
